{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import skimage.io as io\n",
    "import os, sys\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "# Use second GPU -- change if you want to use a first one\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '1'\n",
    "\n",
    "# Add a path to a custom fork of TF-Slim\n",
    "# Get it from here:\n",
    "# https://github.com/warmspringwinds/models/tree/fully_conv_vgg\n",
    "sys.path.append(\"/home/dpakhom1/workspace/my_models/slim/\")\n",
    "\n",
    "# Add path to the cloned library\n",
    "sys.path.append(\"/home/dpakhom1/tf_projects/segmentation/tf-image-segmentation/\")\n",
    "\n",
    "checkpoints_dir = '/home/dpakhom1/checkpoints'\n",
    "log_folder = '/home/dpakhom1/tf_projects/segmentation/log_folder_resnet_101_8s'\n",
    "\n",
    "number_of_epochs = 20\n",
    "\n",
    "slim = tf.contrib.slim\n",
    "resnet_101_v1_checkpoint_path = os.path.join(checkpoints_dir, 'resnet_v1_101.ckpt')\n",
    "\n",
    "from tf_image_segmentation.utils.tf_records import read_tfrecord_and_decode_into_image_annotation_pair_tensors\n",
    "from tf_image_segmentation.models.resnet_v1_101_8s import resnet_v1_101_8s, extract_resnet_v1_101_mapping_without_logits\n",
    "\n",
    "from tf_image_segmentation.utils.pascal_voc import pascal_segmentation_lut\n",
    "\n",
    "from tf_image_segmentation.utils.training import get_valid_logits_and_labels\n",
    "\n",
    "from tf_image_segmentation.utils.augmentation import (distort_randomly_image_color,\n",
    "                                                      flip_randomly_left_right_image_with_annotation,\n",
    "                                                      scale_randomly_image_with_annotation_with_fixed_size_output)\n",
    "\n",
    "from tensorflow.python.ops import control_flow_ops\n",
    "\n",
    "image_train_size = [384, 384]\n",
    "number_of_classes = 21\n",
    "tfrecord_filename = 'pascal_augmented_train.tfrecords'\n",
    "pascal_voc_lut = pascal_segmentation_lut()\n",
    "class_labels = pascal_voc_lut.keys()\n",
    "\n",
    "\n",
    "filename_queue = tf.train.string_input_producer(\n",
    "    [tfrecord_filename], num_epochs=number_of_epochs)\n",
    "\n",
    "image, annotation = read_tfrecord_and_decode_into_image_annotation_pair_tensors(filename_queue)\n",
    "\n",
    "# Various data augmentation stages\n",
    "image, annotation = flip_randomly_left_right_image_with_annotation(image, annotation)\n",
    "\n",
    "# image = distort_randomly_image_color(image)\n",
    "\n",
    "resized_image, resized_annotation = scale_randomly_image_with_annotation_with_fixed_size_output(image, annotation, image_train_size)\n",
    "\n",
    "\n",
    "resized_annotation = tf.squeeze(resized_annotation)\n",
    "\n",
    "image_batch, annotation_batch = tf.train.shuffle_batch( [resized_image, resized_annotation],\n",
    "                                             batch_size=1,\n",
    "                                             capacity=3000,\n",
    "                                             num_threads=2,\n",
    "                                             min_after_dequeue=1000)\n",
    "\n",
    "# is_training=False here means that we fix the mean and variance\n",
    "# of original model that was trained on Imagenet.\n",
    "# https://github.com/tensorflow/tensorflow/issues/1122\n",
    "upsampled_logits_batch, resnet_v1_101_variables_mapping = resnet_v1_101_8s(image_batch_tensor=image_batch,\n",
    "                                                           number_of_classes=number_of_classes,\n",
    "                                                           is_training=False)\n",
    "\n",
    "\n",
    "valid_labels_batch_tensor, valid_logits_batch_tensor = get_valid_logits_and_labels(annotation_batch_tensor=annotation_batch,\n",
    "                                                                                     logits_batch_tensor=upsampled_logits_batch,\n",
    "                                                                                    class_labels=class_labels)\n",
    "\n",
    "\n",
    "\n",
    "cross_entropies = tf.nn.softmax_cross_entropy_with_logits(logits=valid_logits_batch_tensor,\n",
    "                                                          labels=valid_labels_batch_tensor)\n",
    "\n",
    "# Normalize the cross entropy -- the number of elements\n",
    "# is different during each step due to mask out regions\n",
    "cross_entropy_sum = tf.reduce_mean(cross_entropies)\n",
    "\n",
    "pred = tf.argmax(upsampled_logits_batch, dimension=3)\n",
    "\n",
    "probabilities = tf.nn.softmax(upsampled_logits_batch)\n",
    "\n",
    "\n",
    "with tf.variable_scope(\"adam_vars\"):\n",
    "    train_step = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(cross_entropy_sum)\n",
    "\n",
    "\n",
    "# Variable's initialization functions\n",
    "resnet_v1_101_without_logits_variables_mapping = extract_resnet_v1_101_mapping_without_logits(resnet_v1_101_variables_mapping)\n",
    "\n",
    "\n",
    "init_fn = slim.assign_from_checkpoint_fn(model_path=resnet_101_v1_checkpoint_path,\n",
    "                                         var_list=resnet_v1_101_without_logits_variables_mapping)\n",
    "\n",
    "global_vars_init_op = tf.global_variables_initializer()\n",
    "\n",
    "tf.summary.scalar('cross_entropy_loss', cross_entropy_sum)\n",
    "\n",
    "merged_summary_op = tf.summary.merge_all()\n",
    "\n",
    "summary_string_writer = tf.summary.FileWriter(log_folder)\n",
    "\n",
    "# Create the log folder if doesn't exist yet\n",
    "if not os.path.exists(log_folder):\n",
    "     os.makedirs(log_folder)\n",
    "    \n",
    "#The op for initializing the variables.\n",
    "local_vars_init_op = tf.local_variables_initializer()\n",
    "\n",
    "combined_op = tf.group(local_vars_init_op, global_vars_init_op)\n",
    "\n",
    "# We need this to save only model variables and omit\n",
    "# optimization-related and other variables.\n",
    "model_variables = slim.get_model_variables()\n",
    "saver = tf.train.Saver(model_variables)\n",
    "\n",
    "\n",
    "with tf.Session()  as sess:\n",
    "    \n",
    "    sess.run(combined_op)\n",
    "    init_fn(sess)\n",
    "\n",
    "    coord = tf.train.Coordinator()\n",
    "    threads = tf.train.start_queue_runners(coord=coord)\n",
    "    \n",
    "    # 10 epochs\n",
    "    for i in xrange(11127 * number_of_epochs):\n",
    "    \n",
    "        cross_entropy, summary_string, _ = sess.run([ cross_entropy_sum,\n",
    "                                                      merged_summary_op,\n",
    "                                                      train_step ])\n",
    "        \n",
    "        print(\"Current loss: \" + str(cross_entropy))\n",
    "        \n",
    "        summary_string_writer.add_summary(summary_string, i)\n",
    "        \n",
    "        if i % 11127 == 0:\n",
    "            save_path = saver.save(sess, \"/home/dpakhom1/tf_projects/segmentation/model_resnet_101_8s.ckpt\")\n",
    "            print(\"Model saved in file: %s\" % save_path)\n",
    "            \n",
    "        \n",
    "    coord.request_stop()\n",
    "    coord.join(threads)\n",
    "    \n",
    "    save_path = saver.save(sess, \"/home/dpakhom1/tf_projects/segmentation/model_resnet_101_8s.ckpt\")\n",
    "    print(\"Model saved in file: %s\" % save_path)\n",
    "    \n",
    "summary_string_writer.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
